package cz.cuni.mff.ufal.volk.langr;

import java.io.Serializable;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Vector;
import java.util.Map.Entry;

import cz.cuni.mff.ufal.volk.langr.Utils.Pair;

/**
 * An implementation of the {@link CipherDecoder} interface that can be used for decoding texts encoded
 * using the substitution cipher. It uses the frequency analysis and ngram language models to build
 * the substitution table.
 * 
 * @author Bartłomiej Etenkowski
 */
public class NgramDecoder implements CipherDecoder, Serializable {
  private static final long serialVersionUID = 1837173150497964402L;

  /**
   * Constructs an <tt>NgramDecoder</tt>. 
   * @param model underlying language model
   */
  public NgramDecoder(NgramLanguageModel model) {
    this.model = model;
  }
  
  private final NgramLanguageModel model;
  
  private final Map<Character, Character> substitutionTable = new HashMap<Character, Character>();
  
  /**
   * Returns the read-only substitution table used by the decoder.
   * @return the substitution table
   */
  public Map<Character, Character> getSubstitutionTable() {
    return Collections.unmodifiableMap(substitutionTable);
  }
  
  /**
   * Constructs the substitution table based on a sample encoded text.
   * @param encoded a sample encoded text
   */
  public void train(String encoded) {
    substitutionTable.clear();
    List<Pair<Character, Integer>> encodedFreq = getSortedCounts(encoded);
    List<Pair<Character, Double>> decodedFreq = Utils.getSortedItems(model.getCharacterProbabilities());
    int lastDecoded = decodedFreq.size() - 1;
    for (int i = encodedFreq.size() - 1; i >= 0; i--) {
      double bestProb = Double.NEGATIVE_INFINITY; // log
      char encodedChar = encodedFreq.get(i).key;
      char lastCandidate = 0;
      int selected = -1;
      boolean firstIter = true;
      for (int j = lastDecoded; j >= 0; j--) {
        if (firstIter) {
          if (decodedFreq.get(j) == null)
            lastDecoded--;
          else
            firstIter = false;
        }
        if (decodedFreq.get(j) != null) {
          char candidate = decodedFreq.get(j).key;
          substitutionTable.put(encodedChar, candidate);
          String decoded = decode(encoded);
          double prob = model.logProbability(decoded);
          if (prob > bestProb) {
            lastCandidate = candidate;
            bestProb = prob;
            selected = j;
          } else {
            substitutionTable.put(encodedChar, lastCandidate);
          }
        }
      }
      if (selected >= 0)
        decodedFreq.set(selected, null);
    }
  }

  /**
   * Retrieves counts of particular characters contained by the text and sorts them from the least to the most
   * frequent.
   * @param text
   * @return
   */
  private List<Pair<Character, Integer>> getSortedCounts(String text) {
    Map<Character, Integer> charCounts = new HashMap<Character, Integer>();
    for (int i = 0; i < text.length(); i++) {
      Utils.incrementIntValue(charCounts, text.charAt(i));
    }
    List<Pair<Character, Integer>> countVec = new Vector<Pair<Character,Integer>>(charCounts.size());
    for (Entry<Character, Integer> e : charCounts.entrySet()) {
      countVec.add(new Pair<Character, Integer>(e.getKey(), e.getValue()));
    }
    Collections.sort(countVec);
    return countVec;
  }
  
  /**
   * Decodes a single character.
   * @param encoded character to be decoded
   * @return the decoded character
   */
  public char decode(char encoded) {
    return decode(substitutionTable, encoded);
  }

  @Override
  public String decode(String encoded) {
    return decode(substitutionTable, encoded);
  }
  
  /**
   * Decodes a single character using the given substitution table.
   * @param substitutions the substitution table
   * @param encoded the encoded character
   * @return the decoded character
   */
  public static char decode(Map<Character, Character> substitutions, char encoded) {
    Character decoded = substitutions.get(encoded);
    if (decoded == null)
      return encoded;
    else
      return decoded;
  }
  
  /**
   * Decodes whole text using the given substitution table.
   * @param substitutions the substitution table
   * @param text text to be decoded
   * @return the decoded text
   */
  public static String decode(Map<Character, Character> substitutions, String text) {
    StringBuilder sb = new StringBuilder();
    for (int i = 0; i < text.length(); i++)
      sb.append(decode(substitutions, text.charAt(i)));
    return sb.toString();
  }

  @Override
  public String toString() {
    return String.format("n-gram decoder: %s", (substitutionTable == null ? "NULL" : substitutionTable.toString()));
  }
}
